import hashlib
import os
from typing import List

# 常量配置
FOLDER_A = "/folderA"  # 替换为实际文件夹A路径
FOLDER_B = "/folderB"  # 替换为实际文件夹B路径
COMPARE_CONTENT = False  # True表示比较文件内容，False只比较文件名和大小
OUTPUT_FILE = "identical_files.txt"  # 结果输出文件(设为None则打印到屏幕)
SORT_BY = "name"  # 排序方式: "name"(按名称), "size"(按大小), "path"(按路径长度)


def get_file_hash(filepath):
  """计算文件的MD5哈希值"""
  hash_md5 = hashlib.md5()
  with open(filepath, "rb") as f:
    for chunk in iter(lambda: f.read(4096), b""):
      hash_md5.update(chunk)
  return hash_md5.hexdigest()


def sort_items(items: List[str], parent_path: str) -> List[str]:
  """根据SORT_BY设置对文件和目录进行排序"""
  if SORT_BY == "name":
    return sorted(items)
  elif SORT_BY == "size":
    return sorted(items, key=lambda x: os.path.getsize(os.path.join(parent_path, x)))
  elif SORT_BY == "path":
    return sorted(items, key=lambda x: len(os.path.join(parent_path, x)))
  else:
    return sorted(items)


def compare_files(file_a, file_b):
  """比较两个文件是否相同"""
  if COMPARE_CONTENT:
    # 比较内容(使用哈希提高大文件比较效率)
    return get_file_hash(file_a) == get_file_hash(file_b)
  else:
    # 只比较文件名和大小
    return os.path.getsize(file_a) == os.path.getsize(file_b)


def compare_folders(folder_a, folder_b, relative_path=""):
  """递归比较两个文件夹"""
  _identical_files = []

  try:
    # 获取当前目录下的文件和子目录
    items_a = set(os.listdir(folder_a))
    items_b = set(os.listdir(folder_b))
    common_items = items_a & items_b

    for item in common_items:
      path_a = os.path.join(folder_a, item)
      path_b = os.path.join(folder_b, item)
      rel_path = os.path.join(relative_path, item)

      if os.path.isdir(path_a) and os.path.isdir(path_b):
        # 递归比较子目录
        subdir_files = compare_folders(path_a, path_b, rel_path)
        _identical_files.extend(subdir_files)
      elif os.path.isfile(path_a) and os.path.isfile(path_b):
        # 比较文件
        if compare_files(path_a, path_b):
          _identical_files.append(rel_path)
  except PermissionError:
    print(f"警告: 无权限访问 {folder_a} 或 {folder_b} 中的某些内容")
  except Exception as e:
    print(f"错误: 比较 {folder_a} 和 {folder_b} 时出错 - {str(e)}")

  return _identical_files


def write_results(results, output_file):
  """将结果写入文件或打印到屏幕"""
  # 对最终结果进行排序
  if SORT_BY == "name":
    results.sort()
  elif SORT_BY == "size":
    results.sort(key=lambda x: os.path.getsize(os.path.join(FOLDER_A, x)))
  elif SORT_BY == "path":
    results.sort(key=lambda x: len(x))

  if output_file:
    with open(output_file, "w", encoding="utf-8") as f:
      if results:
        f.write("找到以下相同文件:\n")
        f.write("\n".join(results))
        f.write(f"\n\n总计找到 {len(results)} 个相同文件")
        f.write(f"\n排序方式: {SORT_BY}")
      else:
        f.write("没有找到相同文件")
    print(f"结果已保存到 {output_file}")
  else:
    if results:
      print("找到以下相同文件:")
      for file in results:
        print(f"- {file}")
      print(f"\n总计找到 {len(results)} 个相同文件")
      print(f"排序方式: {SORT_BY}")
    else:
      print("没有找到相同文件")


if __name__ == "__main__":
  # 验证文件夹路径
  if not os.path.isdir(FOLDER_A):
    print(f"错误: 文件夹A路径不存在 - {FOLDER_A}")
    exit(1)
  if not os.path.isdir(FOLDER_B):
    print(f"错误: 文件夹B路径不存在 - {FOLDER_B}")
    exit(2)

  print(f"开始比较文件夹:\nA: {FOLDER_A}\nB: {FOLDER_B}")
  print(f"比较模式: {'内容比较' if COMPARE_CONTENT else '仅比较文件名和大小'}")
  print(f"排序方式: {SORT_BY}")

  # 执行比较
  identical_files = compare_folders(FOLDER_A, FOLDER_B)

  # 输出结果
  write_results(identical_files, OUTPUT_FILE)
